#ifndef __POISSON_H__
#define __POISSON_H__

#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
#include <cmath>
#include <ostream>
#include <valarray>
#include "Element.h"
#include "Quadrature.h"

#define TEMPLATE template <typename Element>

typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::Triplet<double> Tri;
typedef Eigen::VectorXd Vector;

double u(const Point<2, double> &p);
double f(const Point<2, double> &p);

// to help factory
class AbstractPoisson
{
public:
	AbstractPoisson() = default;
	virtual ~AbstractPoisson() = default;
	virtual void SetN(size_t _n) = 0;
	virtual size_t GetN() const = 0;
	virtual void SetDomain(const std::valarray<Point<2> >& _domain) = 0;
	virtual const std::valarray<Point<2> >& GetDomain() const = 0;
	virtual void SetParams(size_t _n, const std::valarray<Point<2> >& _domain) = 0;
	virtual bool IsBoundary(const Point<2> &index) const = 0;
	virtual void Compute() = 0;
	virtual const SpMat& GetMatrix() = 0;
	virtual const Vector& GetRhs() = 0;
	virtual Vector Solve() = 0;
	virtual void SaveToFile(const Vector& solution, const std::string& file_name) = 0;
	virtual Vector Exact() const = 0;
	virtual double SquareError(const Vector& solution) = 0;
	virtual double L2Error(const Vector& solution) = 0;
};



template <typename Element = P1Element>
class Poisson : public AbstractPoisson
{
private:
	size_t n;
	SpMat A;
	Vector rhs;
	std::valarray<Point<2> > domain;
	bool is_compute;

public:
	Poisson() = default;
	Poisson(size_t _n, const std::valarray<Point<2> >& _domain);
	~Poisson() = default;
	void SetN(size_t _n);
	size_t GetN() const;
	void SetDomain(const std::valarray<Point<2> >& _domain);
	const std::valarray<Point<2> >& GetDomain() const;
	void SetParams(size_t _n, const std::valarray<Point<2> >& _domain);
	bool IsBoundary(const Point<2> &index) const;
	void Compute();
	const SpMat& GetMatrix();
	const Vector& GetRhs();
	Vector Solve();
	void SaveToFile(const Vector& solution, const std::string& file_name);
	Vector Exact() const;
	double SquareError(const Vector& solution);
	double L2Error(const Vector& solution);
};

TEMPLATE
Poisson<Element>::Poisson(size_t _n, const std::valarray<Point<2> >& _domain) : n(_n), 
	domain(_domain), is_compute(false), A((n - 1) * (n - 1), (n - 1) * (n - 1)) 
{
	if (domain.size() != 2)
		throw std::length_error("domain should be 2 * 2");
}

TEMPLATE
void Poisson<Element>::SetN(size_t _n)
{
	n = _n;
	is_compute = false;
}
	
TEMPLATE
size_t Poisson<Element>::GetN() const
{
	return n;
}
	
TEMPLATE
void Poisson<Element>::SetDomain(const std::valarray<Point<2> >& _domain)
{
	if (domain.size() != 2)
		throw std::length_error("domain should be 2 * 2");
	domain = _domain;
	is_compute = false;
}
	
TEMPLATE
const std::valarray<Point<2> >& Poisson<Element>::GetDomain() const
{
	return domain;
}

TEMPLATE
void Poisson<Element>::SetParams(size_t _n, const std::valarray<Point<2> >& _domain)
{
	n = _n;
	domain = _domain;
	is_compute = false;
}

TEMPLATE
bool Poisson<Element>::IsBoundary(const Point<2> &index) const
{
	return (index[0] == 0) || (index[0] == n) || (index[1] == 0) || (index[1] == n);
}


TEMPLATE
void Poisson<Element>::Compute()
{
	TriangularQuad quad;
	quad.ReadQuadInfo("../data/triangle.quad", 8);
	const std::valarray<Point<2> > &quad_points = quad.GetQuadPoints();
	const std::valarray<double> &weights = quad.GetQuadWeights();
	size_t n_quad_points = quad.GetNQuadPoints();

	Element element;
	size_t node_num = element.GetNodeNum(n);
	A = SpMat(node_num, node_num);
	rhs = Vector(node_num);
	rhs.setZero(node_num);
	std::vector<Tri> TriList;

	for (size_t ele_index = 0; ele_index < 2 * n * n; ++ele_index) {
		const std::valarray<Point<2> > &dofs_index = element.DofsIndex(ele_index, n);
		for (size_t k = 0; k < element.GetNDofs(); ++k)
			element.SetDof(k, Point<2>({domain[0][0] + (domain[0][1] - domain[0][0]) * dofs_index[k][0] * 1.0 / n, 
										domain[1][0] + (domain[1][1] - domain[1][0]) * dofs_index[k][1] * 1.0 / n}));

		for (size_t l = 0; l < n_quad_points; ++l) {
			double detJ = element.GlobalToLocalJacobiDet(quad_points[l]);

			for (size_t i = 0; i < element.GetNDofs(); ++i) {
				std::valarray<double> grad_i = element.LocalToGlobalJacobiTimes(quad_points[l], element.BasisGradient(i, quad_points[l]));
				double phi_i = element.BasisFunction(i, quad_points[l]);
				bool is_boundary_i = IsBoundary(dofs_index[i]);
				for (size_t j = 0; j < element.GetNDofs(); ++j) {
					std::valarray<double> grad_j = element.LocalToGlobalJacobiTimes(quad_points[l], element.BasisGradient(j, quad_points[l]));
					double grad_i_j = (grad_i[0] * grad_j[0] + grad_i[1] * grad_j[1]) * weights[l] * detJ;
					bool is_boundary_j = IsBoundary(dofs_index[j]);
					if (!is_boundary_i && !is_boundary_j) {
						size_t k_i = element.SerialNumOfIndex(dofs_index[i][0], dofs_index[i][1], n);
						size_t k_j = element.SerialNumOfIndex(dofs_index[j][0], dofs_index[j][1], n);
						TriList.push_back(Tri(k_i, k_j, grad_i_j));
					} else if (!is_boundary_i && is_boundary_j) {
						size_t k_i = element.SerialNumOfIndex(dofs_index[i][0], dofs_index[i][1], n);
						rhs[k_i] -= u(element.GetDof(j)) * grad_i_j;
					}
				}
				if (!is_boundary_i) {
					size_t k_i = element.SerialNumOfIndex(dofs_index[i][0], dofs_index[i][1], n);
					rhs[k_i] += f(element.LocalToGlobal(quad_points[l])) * phi_i * weights[l] * detJ;
				}
			}
		}
	}

	A.setFromTriplets(TriList.begin(), TriList.end());
	A.makeCompressed();
	is_compute = true;
}

TEMPLATE
const SpMat& Poisson<Element>::GetMatrix()
{
	if (!is_compute)
		Compute();
	return A;
}

TEMPLATE
const Vector& Poisson<Element>::GetRhs()
{
	if (!is_compute)
		Compute();
	return rhs;
}

TEMPLATE
Vector Poisson<Element>::Solve()
{
	if (!is_compute)
		Compute();
	Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > solver;
    solver.setTolerance(1e-15);
    solver.compute(A);
	return solver.solve(rhs);
}

TEMPLATE
void Poisson<Element>::SaveToFile(const Vector& solution, const std::string& file_name)
{
	std::fstream output(file_name, std::ios::out);
	Element element;
	double num = sqrt(element.GetNodeNum(n)) + 1.0;
	for (double i = n / num; i <= n - n / num; i += n / num) {
		double x = domain[0][0] + (domain[0][1] - domain[0][0]) * i * 1.0 / n; 
		for (double j = n / num; j <= n - n / num; j += n / num) {
			double y = domain[1][0] + (domain[1][1] - domain[1][0]) * j * 1.0 /n;
			size_t k = element.SerialNumOfIndex(i, j, n);
			output << x << " " << y << " " << solution[k] << std::endl;
		}
	}
	output.close();
}

TEMPLATE
Vector Poisson<Element>::Exact() const
{
	Element element;
	double num = sqrt(element.GetNodeNum(n)) + 1.0;
	Vector exact(element.GetNodeNum(n));
	for (double i = n / num; i <= n - n / num; i += n / num) {
		double x = domain[0][0] + (domain[0][1] - domain[0][0]) * i * 1.0 / n; 
		for (double j = n / num; j <= n - n / num; j += n / num) {
			double y = domain[1][0] + (domain[1][1] - domain[1][0]) * j * 1.0 /n;
			size_t k = element.SerialNumOfIndex(i, j, n);
			exact[k] = u(Point<2>({x, y}));
		}
	}
	return exact;
}

TEMPLATE
double Poisson<Element>::SquareError(const Vector& solution)
{
	Element element;
	double error = 0.0, tmp = 0.0;
	double num = sqrt(element.GetNodeNum(n)) + 1.0;
	for (double i = n / num; i <= n - n / num; i += n / num) {
		double x = domain[0][0] + (domain[0][1] - domain[0][0]) * i * 1.0 / n; 
		for (double j = n / num; j <= n - n / num; j += n / num) {
			double y = domain[1][0] + (domain[1][1] - domain[1][0]) * j * 1.0 /n;
			size_t k = element.SerialNumOfIndex(i, j, n);
			tmp = solution[k] - u(Point<2>({x, y}));
			error += tmp * tmp;
		}
	}
	return error;
}

TEMPLATE
double Poisson<Element>::L2Error(const Vector& solution)
{	
	TriangularQuad quad;
	quad.ReadQuadInfo("../data/triangle.quad", 8);
	const std::valarray<Point<2> > &quad_points = quad.GetQuadPoints();
	const std::valarray<double> &weights = quad.GetQuadWeights();
	size_t n_quad_points = quad.GetNQuadPoints();

	double error = 0.0, tmp = 0.0;
	Element element;
	for (size_t ele_index = 0; ele_index < 2 * n * n; ++ele_index) {
		const std::valarray<Point<2> > &dofs_index = element.DofsIndex(ele_index, n);
		for (size_t k = 0; k < element.GetNDofs(); ++k)
			element.SetDof(k, Point<2>({domain[0][0] + (domain[0][1] - domain[0][0]) * dofs_index[k][0] * 1.0 / n, 
										domain[1][0] + (domain[1][1] - domain[1][0]) * dofs_index[k][1] * 1.0 / n}));

		for (size_t l = 0; l < n_quad_points; ++l) {
			double detJ = element.GlobalToLocalJacobiDet(quad_points[l]);
			
			tmp = u(element.LocalToGlobal(quad_points[l]));
			for (size_t i = 0; i < element.GetNDofs(); ++i) {
				std::valarray<double> grad_i = element.LocalToGlobalJacobiTimes(quad_points[l], element.BasisGradient(i, quad_points[l]));
				double phi_i = element.BasisFunction(i, quad_points[l]);
				bool is_boundary_i = IsBoundary(dofs_index[i]);
				if (!is_boundary_i) {
					size_t k_i = element.SerialNumOfIndex(dofs_index[i][0], dofs_index[i][1], n);
					tmp -= solution[k_i] * phi_i;
				} else {
					tmp -= u(element.GetDof(i)) * phi_i;
				}
			}
			error += tmp * tmp * weights[l] * detJ;
		}
	}
	return sqrt(error);
}


std::ostream& operator<<(std::ostream& os, const std::valarray<double>& v);
double MaxNorm(const Vector& v);

#undef TEMPLATE
#endif
